#include <butil/string_printf.h>
#include <gflags/gflags.h>

#include "index/primary_key_index/primary_key_mem_indexer.h"

namespace hawking {
namespace indexlib {

DEFINE_int32(primary_key_default_pool_bytes, 10 * 1024 * 1024, "default pool bytes");
DEFINE_int32(primary_key_hashmap_init_bytes, 1024, "hashmap init bytes");

Status PrimaryKeyMemIndexer::Init() {
    allocator_.reset(new MMapAllocator());
    pool_.reset(new Pool(allocator_.get(), FLAGS_primary_key_default_pool_bytes));

    if (pool_.get() != nullptr) {
        hash_map_.reset(new HashMap<uint64_t, uint64_t>(
            pool_.get(), FLAGS_primary_key_hashmap_init_bytes));
    } else {
        hash_map_.reset(new HashMap<uint64_t, uint64_t>(
            FLAGS_primary_key_hashmap_init_bytes));
    }

    pk_file_writer_.reset(new SortedPrimaryKeyFileWriter<uint64_t>());
    return Status::OK;
}

Status PrimaryKeyMemIndexer::AddDocument(uint64_t pk, DocId doc) {
    hash_map_->FindAndInsert(pk, doc);
    return Status::OK;
}

Status PrimaryKeyMemIndexer::Lookup(
    const google::protobuf::RepeatedField<int64_t>& pks,
    SegmentContext* segment_ctx) {
    segment_ctx->Reserve(pks.size());
    for (const auto pk : pks) {
        DocId* doc_id_ptr = hash_map_->Find(pk);
        if (!doc_id_ptr) {
            continue;
        }

        segment_ctx->PushDocId(*doc_id_ptr);
        segment_ctx->MutableDocContext()->back().pk = pk;
    }

    return Status::OK;
}

ReturnValue<DocId> PrimaryKeyMemIndexer::Lookup(uint64_t pk) {
    DocId* doc_id_ptr = hash_map_->Find(pk);
    if (!doc_id_ptr) {
        return ReturnValue<DocId>{Status::NotFound, 0};
    }

    return ReturnValue<DocId>{Status::OK, *doc_id_ptr};
}

Status PrimaryKeyMemIndexer::Dump(
    const std::string& pk_dir,
    const std::unordered_map<DocId, DocId>& old2new_map) {
    std::string filepath = butil::string_printf("%s/data_file", pk_dir.data());
    std::shared_ptr<NormalFileWriter> file_writer =
        std::make_shared<NormalFileWriter>(filepath);
    Status ret = file_writer->Init();
    if (ret != Status::OK) {
        return ret;
    }
    
    pk_file_writer_->Init(old2new_map.size(), file_writer, pool_.get());

    auto iter = hash_map_->CreateIterator();
    while (iter.HasNext()) {
        auto& p = iter.Next();
        auto iter = old2new_map.find(p.second);
        if (iter == old2new_map.end()) {
            continue;
        }

        DocId new_docid = iter->second;
        auto status = pk_file_writer_->AddPKPair(p.first, new_docid);
        if (status != Status::OK) {
            return status;
        }
    }

    return pk_file_writer_->Close();
}

// no matter mem segment or further disk segment, doc_id in segment is self-increment.
// some doc is deleted before dump, they are recoreded in delete_indexer.
// now dump, don't dump the deleted doc.
std::unordered_map<DocId, DocId> PrimaryKeyMemIndexer::GenerateOldNewDocMap(
    std::shared_ptr<DeletionMapIndexer> deletion_indexer) const {
    std::unordered_map<DocId, DocId> old2new_doc_map;
    old2new_doc_map.reserve(DocCount() - deletion_indexer->GetDeletedDocCount());
    DocId pos = 0;
    for (DocId old_docid = 0; old_docid < hash_map_->Size(); ++old_docid) {
        if (!deletion_indexer->IsDeleted(old_docid)) {
            old2new_doc_map.emplace(old_docid, pos++);
        }
    }

    return old2new_doc_map;
}

}
}